"""
重试处理组件
负责错误重试逻辑
"""
from typing import Callable, Any, Tuple, Dict, Optional, Type
import time
import logging
import random
from enum import Enum, auto

# 设置日志
logger = logging.getLogger("connllm.retry")

class RetryStrategy(Enum):
	"""重试策略枚举"""
	CONSTANT = auto()  # 固定延迟
	LINEAR = auto()    # 线性增长延迟
	EXPONENTIAL = auto()  # 指数增长延迟
	EXPONENTIAL_JITTER = auto()  # 带抖动的指数增长

class RetryHandler:
	"""重试处理器，负责错误重试逻辑"""
	
	def __init__(self, max_retries: int = 3, 
				 retry_delay: float = 1.0,
				 retry_strategy: str = "exponential"):
		"""
		初始化重试处理器
		
		Args:
			max_retries: 最大重试次数
			retry_delay: 基础重试延迟(秒)
			retry_strategy: 重试策略名称
				- constant: 固定延迟
				- linear: 线性增长
				- exponential: 指数增长
				- exponential_jitter: 带抖动的指数增长
		"""
		self.max_retries = max_retries
		self.retry_delay = retry_delay
		
		# 转换策略名称为枚举
		self.retry_strategy = self._parse_strategy(retry_strategy)
		
		# 初始化错误类型映射（延迟至execute方法）
		self._retryable_errors = None
		self._non_retryable_errors = None
	
	def _parse_strategy(self, strategy_name: str) -> RetryStrategy:
		"""
		解析重试策略名称为枚举值
		
		Args:
			strategy_name: 策略名称
			
		Returns:
			重试策略枚举
		"""
		strategy_map = {
			"constant": RetryStrategy.CONSTANT,
			"linear": RetryStrategy.LINEAR,
			"exponential": RetryStrategy.EXPONENTIAL,
			"exponential_jitter": RetryStrategy.EXPONENTIAL_JITTER
		}
		
		return strategy_map.get(strategy_name.lower(), RetryStrategy.EXPONENTIAL)
	
	def execute(self, func: Callable, args: Tuple = (), kwargs: Optional[Dict[str, Any]] = None) -> Any:
		"""
		执行带重试的操作
		
		Args:
			func: 要执行的函数
			args: 位置参数
			kwargs: 关键字参数
			
		Returns:
			函数执行结果
			
		Raises:
			Exception: 所有重试都失败后抛出异常
		"""
		if kwargs is None:
			kwargs = {}
			
		# 延迟初始化错误类型映射
		if self._retryable_errors is None or self._non_retryable_errors is None:
			self._init_error_types()
		
		last_error = None
		retry_count = 0
		
		for attempt in range(self.max_retries + 1):
			try:
				return func(*args, **kwargs)
			except Exception as e:
				# 转换为标准异常
				error = self._map_error(e)
				last_error = error
				
				# 最后一次尝试失败，抛出异常
				if attempt >= self.max_retries:
					break
					
				# 检查是否应该重试
				if not self.should_retry(error):
					logger.info(f"错误不可重试，放弃后续尝试: {str(error)}")
					break
					
				# 计算重试延迟
				retry_count += 1
				delay = self._calculate_delay(retry_count)
				
				# 记录重试信息
				logger.warning(
					f"操作失败，重试 {retry_count}/{self.max_retries}: {str(error)}, "
					f"等待 {delay:.2f}秒"
				)
				
				# 等待重试
				time.sleep(delay)
		
		# 所有重试都失败
		if last_error:
			logger.error(f"所有重试都失败了: {str(last_error)}")
			raise last_error
			
		# 不应该到达这里
		raise RuntimeError("重试逻辑错误：既没有返回结果也没有抛出异常")
		
	def _init_error_types(self):
		"""初始化可重试和不可重试的错误类型"""
		# 导入标准异常类型
		try:
			from ..core.exceptions import (
				AuthenticationError, QuotaExceededError, 
				RateLimitError, ConnectionError, TimeoutError, 
				ServerError, InvalidRequestError
			)
			
			# 可以重试的错误类型
			self._retryable_errors = (
				ConnectionError,
				TimeoutError, 
				ServerError,
				RateLimitError
			)
			
			# 不重试的错误类型
			self._non_retryable_errors = (
				AuthenticationError,
				QuotaExceededError,
				InvalidRequestError
			)
		except ImportError as e:
			logger.warning(f"导入异常类型失败，使用默认设置: {str(e)}")
			# 使用空元组作为默认值
			self._retryable_errors = ()
			self._non_retryable_errors = ()
	
	def should_retry(self, error: Exception) -> bool:
		"""
		判断是否应该重试
		
		Args:
			error: 异常
			
		Returns:
			是否重试
		"""
		# 检查是否是明确可重试的错误类型
		if isinstance(error, self._retryable_errors):
			return True
			
		# 检查是否是明确不可重试的错误类型
		if isinstance(error, self._non_retryable_errors):
			return False
			
		# 如果是未知类型，通过错误消息判断
		error_str = str(error).lower()
		
		# 常见可重试错误模式
		retryable_patterns = [
			'timeout', '超时',
			'connection', '连接',
			'network', '网络',
			'rate limit', '速率限制',
			'too many requests', '请求过多',
			'server error', '服务器错误',
			'5xx'
		]
		
		# 明确不应该重试的错误模式
		non_retryable_patterns = [
			'authentication', '认证',
			'authorization', '授权',
			'permission', '权限',
			'quota', '配额',
			'invalid request', '无效请求',
			'validation', '验证',
			'bad request', '请求错误'
		]
		
		# 检查错误消息是否匹配可重试模式
		for pattern in retryable_patterns:
			if pattern in error_str:
				return True
				
		# 检查错误消息是否匹配不可重试模式
		for pattern in non_retryable_patterns:
			if pattern in error_str:
				return False
				
		# 默认情况下，其他错误不重试
		return False
		
	def _calculate_delay(self, attempt: int) -> float:
		"""
		计算重试延迟
		
		Args:
			attempt: 当前重试次数(从1开始)
			
		Returns:
			延迟时间(秒)
		"""
		if self.retry_strategy == RetryStrategy.CONSTANT:
			return self.retry_delay
			
		elif self.retry_strategy == RetryStrategy.LINEAR:
			return self.retry_delay * attempt
			
		elif self.retry_strategy == RetryStrategy.EXPONENTIAL:
			# 指数退避，基础延迟 * 2^(attempt-1)
			return self.retry_delay * (2 ** (attempt - 1))
			
		elif self.retry_strategy == RetryStrategy.EXPONENTIAL_JITTER:
			# 带抖动的指数退避
			base_delay = self.retry_delay * (2 ** (attempt - 1))
			# 添加0-1秒的随机抖动
			jitter = random.uniform(0, 1)
			return base_delay + jitter
			
		# 默认为指数退避
		return self.retry_delay * (2 ** (attempt - 1))
		
	def _map_error(self, error: Exception) -> Exception:
		"""
		将异常映射为标准异常
		
		Args:
			error: 原始异常
			
		Returns:
			映射后的标准异常
		"""
		# 如果已经是定义好的异常，直接返回
		if self._is_known_error_type(error):
			return error
			
		# 尝试使用异常映射函数
		try:
			from ..core.exceptions import map_provider_error, LLMException
			
			# 如果已经是标准异常，直接返回
			if isinstance(error, LLMException):
				return error
				
			# 否则尝试映射
			provider_type = getattr(self, '_provider_type', None)
			model = getattr(self, '_model', None)
			
			return map_provider_error(provider_type, error, model)
			
		except ImportError:
			# 如果导入失败，直接返回原始异常
			logger.warning(f"无法导入异常映射函数，直接返回原始异常: {str(error)}")
			return error
			
	def _is_known_error_type(self, error: Exception) -> bool:
		"""
		检查是否为已知的错误类型
		
		Args:
			error: 异常
			
		Returns:
			是否为已知错误类型
		"""
		return isinstance(error, self._retryable_errors) or isinstance(error, self._non_retryable_errors)
